# 导入所需模块
from llama_index.core import (
    VectorStoreIndex,
    SimpleDirectoryReader,
    Settings,
    PromptTemplate,
)
from llama_index.llms.ollama import Ollama
from llama_index.embeddings.ollama import OllamaEmbedding
import os

# 1. 设置 Ollama 模型（LLM 和 Embedding）
# 使用 llama3 作为生成模型
Settings.llm = Ollama(model="qwen3:8b", request_timeout=300.0, base_url="http://localhost:11434")

# 使用 Ollama 提供的嵌入模型（如 nomic-embed-text，也可用 mxbai-embed-large）
Settings.embed_model = OllamaEmbedding(
    model_name="nomic-embed-text:v1.5",  # 推荐：高质量通用嵌入
    base_url="http://localhost:11434",
    ollama_additional_kwargs={"keep_alive": 1}  # 模型常驻内存，加快响应
)

# 2. 加载文档（将你的文档放入 'data/' 目录）
# 支持 .pdf, .txt, .docx, .csv 等
documents = SimpleDirectoryReader("data").load_data()

# 3. 构建索引（自动分块、向量化并存入内存）
index = VectorStoreIndex.from_documents(documents)

# 4. 创建查询引擎
query_engine = index.as_query_engine(streaming=True)  # streaming=True 可实现流式输出

# 5. 开始问答
while True:
    question = input("\n问一个问题 (输入 'quit' 退出): ")
    if question.lower() == "quit":
        break
    response = query_engine.query(question)
    print("\n回答: ")
    response.print_response_stream()  # 流式打印回答